package com.wxh.rsa;

import sun.misc.BASE64Encoder;

import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;

/**
 * @ClassName: KeyPair
 * @Description:
 * @Author wxh
 * @Date: 2020/5/9 17:39
 * @Version V1.0.0
 * @Since 1.8
 */
public class KeyPair {
    private java.security.KeyPair _rsa;
    private KeyFormat _format;
    private String _private;
    private String _public;

    public KeyFormat getFormat() {
        return this._format;
    }

    KeyPair(KeyFormat format) throws NoSuchAlgorithmException {
        KeyPairGenerator kpg = KeyPairGenerator.getInstance("RSA");
        kpg.initialize(1024);
        this._rsa = kpg.genKeyPair();
        this._format = format;
    }

    private KeyPair(java.security.KeyPair rsa, KeyFormat format) {
        this._rsa = rsa;
        this._format = format;
    }

    public String getPrivateKey() {
        if (this._private == null) {
            switch (this._format) {
                case ASN:
                    this._private = _toASNPrivateKey();
                    break;
                case PEM:
                    this._private = _toXMLPrivateKey();
                    break;
                case XML:
                    this._private = _toPEMPrivateKey();
                    break;
                default:
                    this._private = _toXMLPrivateKey();
            }
        }
        return this._private;
    }

    public String getPublicKey() {
        if (this._public == null) {
            switch (this._format) {
                case ASN:
                    this._public = _toASNPublicKey();
                    break;
                case PEM:
                    this._public = _toXMLPublicKey();
                    break;
                case XML:
                    this._public = _toPEMPublicKey();
                    break;
                default:
                    this._public = _toXMLPublicKey();
            }
        }
        return this._public;
    }

    public KeyPair toASNKeyPair() {
        return new KeyPair(this._rsa, KeyFormat.ASN);
    }

    public KeyPair toXMLKeyPair() {
        return new KeyPair(this._rsa, KeyFormat.XML);
    }

    public KeyPair toPEMKeyPair() {
        return new KeyPair(this._rsa, KeyFormat.PEM);
    }

    private String _toASNPublicKey() {
        return new BASE64Encoder().encodeBuffer(this._rsa.getPublic().getEncoded());
    }

    private String _toASNPrivateKey() {
        return new BASE64Encoder().encodeBuffer(this._rsa.getPrivate().getEncoded());
    }

    private String _toXMLPublicKey() {
        return XmlKeyBuilder.publicKeyToXML(this._rsa.getPublic());
    }

    private String _toXMLPrivateKey() {
        return XmlKeyBuilder.privateKeyToXML(this._rsa.getPrivate());
    }

    private String _toPEMPublicKey() {
        String publicKey = _toASNPublicKey();
        StringBuilder sb = new StringBuilder();
        sb.append("-----BEGIN PUBLIC KEY-----\r\n");
        int i = 0;
        while (i + 64 < publicKey.length()) {
            sb.append(publicKey.substring(i, i + 64) + "\r\n");
            i += 64;
        }
        sb.append(publicKey.substring(i, publicKey.length()) + "\r\n");
        sb.append("-----END PUBLIC KEY-----\r\n");

        return sb.toString();
    }

    private String _toPEMPrivateKey() {
        String privateKey = _toASNPrivateKey();
        StringBuilder sb = new StringBuilder();
        sb.append("-----BEGIN PRIVATE KEY-----\r\n");
        int i = 0;
        while (i + 64 < privateKey.length()) {
            sb.append(privateKey.substring(i, i + 64) + "\r\n");
            i += 64;
        }
        sb.append(privateKey.substring(i, privateKey.length()) + "\r\n");
        sb.append("-----END PRIVATE KEY-----\r\n");

        return sb.toString();
    }
}
